from modelscope import AutoModelForCausalLM, AutoTokenizer
import traceback

# 修改为实际存在的模型路径
model_name = "/mnt/codes/qwen3/deploy/models/Qwen3-4B"

tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(
    model_name, 
    torch_dtype="auto",  # 指定模型权重数据类型
    device_map="auto"  # 指定模型权重加载设备
)

# - 创建消息message
prompt = "给出一个关于大语言模型的简短介绍。"
messages = [
    {"role": "user", "content": prompt}
]

# - 词嵌入过程
text = tokenizer.apply_chat_template(
    messages,
    tokenize=False,  # 是否对输出进行分词，False表示输出为字符串格式
    add_generation_prompt=True,  # 是否在消息末尾添加生成提示，通常用于生成任务
    enable_thinking=True  # 是否启用"思考"模式，开启后模型会输出<|thinking|>标签内的思考内容
)
model_inputs = tokenizer([text], return_tensors="pt").to(model.device)

# 创建并回复

# 进行文本补全
generated_ids = model.generate(
    **model_inputs,
    max_new_tokens=512  # 减少最大新token数量以提高速度
)
output_ids = generated_ids[0][len(model_inputs.input_ids[0]):].tolist() 

# 结果将首先显示思考内容，在<thinking>标签中，然后是实际响应
print(tokenizer.decode(output_ids, skip_special_tokens=True))